'''
    生成地面单位的攻击策略,通过双方的坐标位置，攻击状态判断是否攻击
'''
import copy
import numpy as np


class enemy_Strategy:
    def __init__(self,f_number, e_number, friends, enemies, dt):
        self.n_friend = f_number        # 地面单位个数
        self.n_enemy = e_number     # 空中单位个数
        self.friends = friends      # 地面单位的友方
        self.enemies = enemies      # 地面单位的敌方--无人机
        # self.att_target_e = att_target_e
        # self.Dis = [[],[],[],[]]        # 用于存放地面第i个单位对第j架无人机的距离
        self.dt = dt

    '''获取双方单位的存活状态'''
    def live(self):
        live = [[],[]]      # 前一个存放地面单位存活状态，后一个存放地面单位存活状态
        for friend in self.friends:
            live[0].append(friend.live)
        for enemy in self.enemies:
            live[1].append(enemy.live)
        return live

    '''获取地面单位的冷却状态'''
    def CD(self):
        CD = []
        for friend in self.friends:
            attack = friend.Attack(self.dt)
            CD.append(attack)
        return CD

    '''形成可攻击单位矩阵'''
    def att_Ok(self, live, CD):
        att_Ok = np.zeros((self.n_friend, self.n_enemy))      # 用于存放可攻击单位矩阵,初始化全部不可攻击
        for i in range(self.n_friend):
            if live[0][i] == 1 and CD[i]:
                for j in range(self.n_enemy):
                    if live[1][j] == 1:
                        att_Ok[i][j] = 1
                    else:
                        att_Ok[i][j] = 0
            else:
                att_Ok[i] = np.zeros(self.n_enemy)
        return att_Ok

    '''获取双方坐标，计算单位间的距离，并返回距离矩阵'''
    def distance(self):
        dis = []
        i = 0
        for friend in self.friends:     # friends 是一个列表，逐个取出单位的坐标
            dis.append([])
            for enemy in self.enemies:
                distance = friend.distance_xyz(friend.Position, enemy.Position)
                dist = copy.deepcopy(distance)
                dis[i].append(dist)
            i += 1
        return dis


    ##############################################修改到此############################################################

    '''形成可攻击单位距离矩阵，不可攻击单位的距离设为一个极大值'''
    def att_Dis(self, att_Ok, dis):
        for i in range(self.n_friend):
            for j in range(self.n_enemy):
                if att_Ok[i][j] == 0:
                    dis[i][j] = 99999
        return dis

    '''最小元素索引，取出每一个地面单位对空中单位可攻击单位距离矩阵中距离最小的元素的索引坐标'''
    def index1(self, att_Dis):
        index = [[],[]]     # 用于存放索引
        i = 0
        for dis in att_Dis:
            j = np.argmin(dis)
            index[0].append(i)      # 存储地面单位序号
            index[1].append(j)      # 存储空中单位序号
            i += 1
        return index

    '''次小元素索引，取出每一个地面单位对空中单位可攻击单位距离矩阵中距离次小的元素的索引坐标'''
    def index2(self, att_Dis):
        index = [[], []]  # 用于存放索引
        value_max = 0
        i = 0
        # 遍历攻击单位距离矩阵att_Dis，对某个地面单位的距离数组dis进行排序，找到次小元素读取其索引
        for dis in att_Dis:
            dis_copy = copy.deepcopy(dis)
            dis_copy.sort()
            value = dis_copy[1]
            j = dis.index(value)
            index[0].append(i)  # 存储地面单位序号
            index[1].append(j)  # 存储空中单位序号
            i += 1
        return index

    '''确定地面单位的攻击策略'''
    def strategy(self, live, index1, index2, CD):
        strategy = np.zeros((self.n_friend, self.n_enemy))  # 初始化策略矩阵，默认所有地面单位都不发动攻击
        j = 0
        sum = 0
        # 对存活单位的攻击策略进行更改
        if self.n_enemy >= 2:
            for i in live[0]:
                if i == 1 and CD[j]:
                    line = index1[0][j]
                    row = index1[1][j]
                    # 判断所选敌方单位是否已有多于2个己方单位选的攻击,若大于两个选择次小距离单位进行攻击
                    for k in strategy:
                        if k[row] == 1:
                            sum += 1
                    if sum >= 2:
                        line_next = index2[0][j]
                        row_next = index2[1][j]
                        if live[1][row_next] == 1:
                            strategy[line_next][row_next] = 1
                    else:
                        if live[1][row] == 1:
                            strategy[line][row] = 1
                j += 1
        else:
            for i in live[0]:
                if i == 1 and CD[j]:
                    line = index1[0][j]
                    row = index1[1][j]
                    if live[1][row] == 1:
                        strategy[line][row] = 1
                j += 1
        return strategy

    '''主函数，调用其他函数，完成地面单位的攻击策略的生成任务'''
    def main(self):
        live = self.live()
        CD = self.CD()
        att_Ok = self.att_Ok(live, CD)
        dis = self.distance()
        att_Dis = self.att_Dis(att_Ok, dis)
        index1 = self.index1(att_Dis)
        if self.n_enemy >= 2:
            index2 = self.index2(att_Dis)
        else:
            index2 = [0, 0]
        strategy = self.strategy(live, index1, index2, CD)

        return strategy
